`timescale 1ns / 1ps

module tb_aru_reduce_max_min_2_to_1;

    // BF16类型定义
    typedef struct packed {
        logic       sign;
        logic [7:0] exp;
        logic [6:0] mant;
    } bf16_t;

    // 信号声明
    logic  clk;
    logic  rst_n;
    logic  en;
    logic  is_max;
    logic  psum_en;
    logic  end_reduce;
    logic  dat_in_vld  [1:0];
    bf16_t dat_in      [1:0];
    bf16_t dat_out;
    logic  dat_out_vld;

    // 实例化DUT
    aru_reduce_max_min_2_to_1 dut (
        .clk(clk),
        .rst_n(rst_n),
        .en(en),
        .is_max(is_max),
        .psum_en(psum_en),
        .end_reduce(end_reduce),
        .dat_in_vld(dat_in_vld),
        .dat_in(dat_in),
        .dat_out(dat_out),
        .dat_out_vld(dat_out_vld)
    );

    // 时钟生成
    initial begin
        clk = 0;
        forever #5 clk = ~clk;
    end

    // BF16转换函数
    function bf16_t real_to_bf16(real value);
        bf16_t result;
        real   abs_val;
        int    exp_val;
        real   mant_val;

        if (value < 0) begin
            result.sign = 1'b1;
            abs_val     = -value;
        end else begin
            result.sign = 1'b0;
            abs_val     = value;
        end

        if (abs_val == 0.0) begin
            result.exp  = 8'h00;
            result.mant = 7'h00;
            return result;
        end

        exp_val  = 0;
        mant_val = abs_val;

        if (mant_val >= 2.0) begin
            while (mant_val >= 2.0 && exp_val < 128) begin
                mant_val = mant_val / 2.0;
                exp_val  = exp_val + 1;
            end
        end else if (mant_val < 1.0) begin
            while (mant_val < 1.0 && exp_val > -126) begin
                mant_val = mant_val * 2.0;
                exp_val  = exp_val - 1;
            end
        end

        if (exp_val > 127) begin
            result.exp  = 8'hFF;
            result.mant = 7'h00;
            return result;
        end else if (exp_val < -126) begin
            result.exp  = 8'h00;
            result.mant = 7'h00;
            return result;
        end

        result.exp  = exp_val + 127;
        mant_val    = mant_val - 1.0;
        mant_val    = mant_val * 128.0;
        result.mant = $rtoi(mant_val);

        return result;
    endfunction

    function real bf16_to_real(bf16_t bf16);
        real result;
        real mant_val;
        int  exp_val;

        if (bf16.exp == 8'hFF) begin
            if (bf16.mant == 0) begin
                result = bf16.sign ? -1.0 / 0.0 : 1.0 / 0.0;
            end else begin
                result = 0.0 / 0.0;
            end
            return result;
        end

        if (bf16.exp == 0 && bf16.mant == 0) begin
            return 0.0;
        end

        exp_val  = bf16.exp - 127;
        mant_val = 1.0 + (real'(bf16.mant) / 128.0);

        result   = mant_val;
        if (exp_val > 0) begin
            for (int i = 0; i < exp_val; i++) begin
                result = result * 2.0;
            end
        end else if (exp_val < 0) begin
            for (int i = 0; i < -exp_val; i++) begin
                result = result / 2.0;
            end
        end

        if (bf16.sign) result = -result;

        return result;
    endfunction

    // 特殊值定义
    localparam bf16_t POS_INF_BF16 = '{sign: 1'b0, exp: 8'hFF, mant: 7'h00};
    localparam bf16_t NEG_INF_BF16 = '{sign: 1'b1, exp: 8'hFF, mant: 7'h00};
    localparam bf16_t NAN_BF16 = '{sign: 1'b0, exp: 8'hFF, mant: 7'h40};
    localparam bf16_t POS_ZERO_BF16 = '{sign: 1'b0, exp: 8'h00, mant: 7'h00};
    localparam bf16_t NEG_ZERO_BF16 = '{sign: 1'b1, exp: 8'h00, mant: 7'h00};

    // 期望值队列
    typedef struct {
        real expected;
        string test_name;
        real tolerance;
        int timestamp;
        logic should_check;
        logic is_special;
    } expected_result_t;

    expected_result_t expected_queue      [$];
    int               input_timestamp = 0;

    // 发送数据任务
    task send_data(input logic [0:1] vld, input real d0, d1, input logic is_max_op, input logic psum,
                   input logic end_red, input string test_name = "", input real expected = 0.0,
                   input real tolerance = 0.01);
        bf16_t            temp[2];
        expected_result_t exp;

        @(posedge clk);
        en            <= 1'b1;
        is_max        <= is_max_op;
        psum_en       <= psum;
        end_reduce    <= end_red;
        dat_in_vld[0] <= vld[0];
        dat_in_vld[1] <= vld[1];

        temp[0] = real_to_bf16(d0);
        temp[1] = real_to_bf16(d1);

        if (vld[0]) dat_in[0] <= temp[0];
        if (vld[1]) dat_in[1] <= temp[1];

        $display("[%0t] Send #%0d: %s vld=%02b %s psum=%b end=%b", $time, input_timestamp, test_name, vld,
                 is_max_op ? "MAX" : "MIN", psum, end_red);
        if (vld[0]) $display("  [0]=%f", d0);
        if (vld[1]) $display("  [1]=%f", d1);

        exp.expected     = expected;
        exp.test_name    = test_name;
        exp.tolerance    = tolerance;
        exp.timestamp    = input_timestamp;
        exp.should_check = end_red;
        exp.is_special   = 0;

        if (test_name != "" && end_red) begin
            expected_queue.push_back(exp);
            $display("  -> Added to queue (queue_size=%0d)", expected_queue.size());
        end

        input_timestamp++;
    endtask

    // 发送特殊值数据
    task send_special_data(input logic [1:0] vld, input bf16_t d0, d1, input logic is_max_op, input logic psum,
                           input logic end_red, input string test_name = "");
        expected_result_t exp;

        @(posedge clk);
        en            <= 1'b1;
        is_max        <= is_max_op;
        psum_en       <= psum;
        end_reduce    <= end_red;
        dat_in_vld[0] <= vld[0];
        dat_in_vld[1] <= vld[1];
        if (vld[0]) dat_in[0] <= d0;
        if (vld[1]) dat_in[1] <= d1;

        $display("[%0t] Send Special #%0d: %s vld=%02b %s psum=%b end=%b", $time, input_timestamp, test_name, vld,
                 is_max_op ? "MAX" : "MIN", psum, end_red);

        exp.expected     = 0.0;
        exp.test_name    = test_name;
        exp.tolerance    = 0.0;
        exp.timestamp    = input_timestamp;
        exp.should_check = 0;
        exp.is_special   = 1;

        if (test_name != "") begin
            expected_queue.push_back(exp);
        end

        input_timestamp++;
    endtask

    // 输出监控进程
    int output_count = 0;
    always @(posedge clk) begin
        if (dat_out_vld) begin
            real              actual;
            expected_result_t exp;
            real              rel_error;

            actual = bf16_to_real(dat_out);

            $display("[%0t] === OUTPUT #%0d === 0x%04h = %f", $time, output_count, dat_out, actual);

            if (expected_queue.size() > 0) begin
                exp = expected_queue.pop_front();

                $display("  Checking against: %s (should_check=%0b)", exp.test_name, exp.should_check);

                if (exp.should_check && !exp.is_special) begin
                    if (exp.expected != 0.0) begin
                        rel_error = $abs((actual - exp.expected) / exp.expected);
                    end else begin
                        rel_error = $abs(actual - exp.expected);
                    end

                    if (exp.expected != exp.expected && actual != actual) begin
                        $display("  [PASS] %s: Both NaN", exp.test_name);
                    end else if (rel_error < exp.tolerance) begin
                        $display("  [PASS] %s: Expected=%f Got=%f Error=%f%%", exp.test_name, exp.expected, actual,
                                 rel_error * 100);
                    end else begin
                        $display("  [FAIL] %s: Expected=%f Got=%f Error=%f%%", exp.test_name, exp.expected, actual,
                                 rel_error * 100);
                    end
                end else begin
                    $display("  [INFO] %s: Output=%f (no check)", exp.test_name, actual);
                end
            end else begin
                $display("  [WARN] No expected value in queue");
            end

            output_count++;
            $display("====== OUTPUT END ======\n");
        end
    end

    // 主测试流程
    initial begin
        $display("========================================");
        $display("Starting BF16 2-to-1 Max/Min Reduce Testbench");
        $display("========================================");

        // 初始化
        rst_n      = 0;
        en         = 0;
        is_max     = 0;
        psum_en    = 0;
        end_reduce = 0;
        for (int i = 0; i < 2; i++) begin
            dat_in_vld[i] = 0;
            dat_in[i]     = '0;
        end

        #20;
        rst_n = 1;
        #10;

        $display("\n=== BASIC MAX TESTS ===");

        send_data(2'b11, 1.0, 2.0, 1'b1, 1'b0, 1'b1, "Test1: Basic MAX", 2.0);
        send_data(2'b11, -1.0, -2.0, 1'b1, 1'b0, 1'b1, "Test2: Negative MAX", -1.0);
        send_data(2'b11, 5.0, -3.0, 1'b1, 1'b0, 1'b1, "Test3: Mixed MAX", 5.0);
        send_data(2'b10, 10.0, 99.0, 1'b1, 1'b0, 1'b1, "Test4: Partial[10] MAX", 10.0);
        send_data(2'b01, 99.0, 15.0, 1'b1, 1'b0, 1'b1, "Test5: Partial[01] MAX", 15.0);

        $display("\n=== BASIC MIN TESTS ===");

        send_data(2'b11, 1.0, 2.0, 1'b0, 1'b0, 1'b1, "Test6: Basic MIN", 1.0);
        send_data(2'b11, -1.0, -2.0, 1'b0, 1'b0, 1'b1, "Test7: Negative MIN", -2.0);
        send_data(2'b11, 5.0, -3.0, 1'b0, 1'b0, 1'b1, "Test8: Mixed MIN", -3.0);
        send_data(2'b10, 10.0, 99.0, 1'b0, 1'b0, 1'b1, "Test9: Partial[10] MIN", 10.0);
        send_data(2'b01, 99.0, 15.0, 1'b0, 1'b0, 1'b1, "Test10: Partial[01] MIN", 15.0);

        $display("\n=== PSUM MAX TESTS ===");

        send_data(2'b11, 1.0, 2.0, 1'b1, 1'b1, 1'b0, "PSUM MAX R1");
        send_data(2'b11, 3.0, 4.0, 1'b1, 1'b1, 1'b0, "PSUM MAX R2");
        send_data(2'b11, 5.0, 6.0, 1'b1, 1'b1, 1'b1, "Test11: PSUM MAX", 6.0);

        send_data(2'b11, -10.0, -5.0, 1'b1, 1'b1, 1'b0, "PSUM MAX NEG R1");
        send_data(2'b11, -20.0, -15.0, 1'b1, 1'b1, 1'b0, "PSUM MAX NEG R2");
        send_data(2'b11, -30.0, -25.0, 1'b1, 1'b1, 1'b1, "Test12: PSUM MAX NEG", -5.0);

        $display("\n=== PSUM MIN TESTS ===");

        send_data(2'b11, 1.0, 2.0, 1'b0, 1'b1, 1'b0, "PSUM MIN R1");
        send_data(2'b11, 3.0, 4.0, 1'b0, 1'b1, 1'b0, "PSUM MIN R2");
        send_data(2'b11, 5.0, 6.0, 1'b0, 1'b1, 1'b1, "Test13: PSUM MIN", 1.0);

        send_data(2'b11, -10.0, -5.0, 1'b0, 1'b1, 1'b0, "PSUM MIN NEG R1");
        send_data(2'b11, -20.0, -15.0, 1'b0, 1'b1, 1'b0, "PSUM MIN NEG R2");
        send_data(2'b11, -30.0, -25.0, 1'b0, 1'b1, 1'b1, "Test14: PSUM MIN NEG", -30.0);

        $display("\n=== LARGE PSUM TESTS ===");

        // 16轮MAX累积
        send_data(2'b11, 1.0, 2.0, 1'b1, 1'b1, 1'b0, "PSUM-16 MAX R1");
        send_data(2'b11, 3.0, 4.0, 1'b1, 1'b1, 1'b0, "PSUM-16 MAX R2");
        send_data(2'b11, 5.0, 6.0, 1'b1, 1'b1, 1'b0, "PSUM-16 MAX R3");
        send_data(2'b11, 7.0, 8.0, 1'b1, 1'b1, 1'b0, "PSUM-16 MAX R4");
        send_data(2'b11, 9.0, 10.0, 1'b1, 1'b1, 1'b0, "PSUM-16 MAX R5");
        send_data(2'b11, 11.0, 12.0, 1'b1, 1'b1, 1'b0, "PSUM-16 MAX R6");
        send_data(2'b11, 13.0, 14.0, 1'b1, 1'b1, 1'b0, "PSUM-16 MAX R7");
        send_data(2'b11, 3.0, 4.0, 1'b1, 1'b1, 1'b1, "Test15: PSUM-16 MAX", 14.0);

        // 16轮MIN累积
        send_data(2'b11, 16.0, 15.0, 1'b0, 1'b1, 1'b0, "PSUM-16 MIN R1");
        send_data(2'b11, 14.0, 13.0, 1'b0, 1'b1, 1'b0, "PSUM-16 MIN R2");
        send_data(2'b11, 2.0, 1.0, 1'b0, 1'b1, 1'b0, "PSUM-16 MIN R3");
        send_data(2'b11, 10.0, 9.0, 1'b0, 1'b1, 1'b0, "PSUM-16 MIN R4");
        send_data(2'b11, 8.0, 7.0, 1'b0, 1'b1, 1'b0, "PSUM-16 MIN R5");
        send_data(2'b11, 6.0, 5.0, 1'b0, 1'b1, 1'b0, "PSUM-16 MIN R6");
        send_data(2'b11, 4.0, 3.0, 1'b0, 1'b1, 1'b0, "PSUM-16 MIN R7");
        send_data(2'b11, 6.0, 5.0, 1'b0, 1'b1, 1'b1, "Test16: PSUM-16 MIN", 1.0);

        $display("\n=== SPECIAL VALUE TESTS ===");

        send_special_data(2'b11, POS_INF_BF16, real_to_bf16(5.0), 1'b1, 1'b0, 1'b1, "Test17: +INF MAX");
        send_special_data(2'b11, NEG_INF_BF16, real_to_bf16(5.0), 1'b1, 1'b0, 1'b1, "Test18: -INF MAX");
        send_special_data(2'b11, NAN_BF16, real_to_bf16(5.0), 1'b1, 1'b0, 1'b1, "Test19: NaN MAX");

        send_special_data(2'b11, POS_INF_BF16, real_to_bf16(5.0), 1'b0, 1'b0, 1'b1, "Test20: +INF MIN");
        send_special_data(2'b11, NEG_INF_BF16, real_to_bf16(5.0), 1'b0, 1'b0, 1'b1, "Test21: -INF MIN");
        send_special_data(2'b11, NAN_BF16, real_to_bf16(5.0), 1'b0, 1'b0, 1'b1, "Test22: NaN MIN");

        $display("\n=== EDGE CASE TESTS ===");

        send_data(2'b11, 0.0, 0.0, 1'b1, 1'b0, 1'b1, "Test23: Zero MAX", 0.0);
        send_data(2'b11, 0.0, 0.0, 1'b0, 1'b0, 1'b1, "Test24: Zero MIN", 0.0);
        send_data(2'b11, 3.14, 3.14, 1'b1, 1'b0, 1'b1, "Test25: Same MAX", 3.14, 0.01);
        send_data(2'b11, 3.14, 3.14, 1'b0, 1'b0, 1'b1, "Test26: Same MIN", 3.14, 0.01);
        // send_data(2'b00, 99.0, 99.0, 1'b1, 1'b0, 1'b1, "Test27: All Invalid MAX", 0.0); //理想结果是inf
        // send_data(2'b00, 99.0, 99.0, 1'b0, 1'b0, 1'b1, "Test28: All Invalid MIN", 0.0); //理想结果是inf

        // 等待所有输出完成
        repeat (20) @(posedge clk);

        $display("\n========================================");
        $display("Testbench Completed");
        $display("Total outputs: %0d", output_count);
        $display("Remaining in queue: %0d", expected_queue.size());
        $display("========================================");
        $finish;
    end

    initial begin
        $fsdbDumpfile("./aru_reduce_max_min_2_to_1.fsdb");
        $fsdbDumpvars("+all");
    end

    // 超时保护
    initial begin
        #50000;
        $display("ERROR: Testbench timeout!");
        $finish;
    end

endmodule
